/*
 * Copyright (c) 2021 Tero Hänninen
 *
 * Usage:
 *      llshow file.ll
 *      llshow file.ll funcname
 */

#setup[linux]
extern(C) {
    //#extern #gshared FILE* stdin;
    #extern #gshared FILE* stdout;
    #extern #gshared FILE* stderr;

    define STDOUT_FILENO = 1;

    enum {
        SEEK_SET,   // relative to beginning
        SEEK_CUR,   // relative to current position
        SEEK_END,   // relative to end
    }

    enum FILE;
    alias c_long = i64;

    void*   malloc(usz size);
    void*   realloc(const void* ptr, usz size);
    void    free(const void* ptr);
    void*   memset(void* s, int c, usz n);
    void*   memcpy(void* dest, const void* src, usz n);
    int     memcmp(const void* s1, const void* s2, usz n);
    usz     strlen(const ubyte* s);
    int     fprintf(FILE* stream, const ubyte* format, ...);
    int     printf(const ubyte* fmt, ...);
    int     isatty(int fd);

    FILE*   fopen(const ubyte* filename, const ubyte* mode);
    usz     fread(void* ptr, usz size, usz nmemb, FILE* stream);
    int     fseek(FILE* stream, c_long offset, int whence);
    c_long  ftell(FILE* stream);
    int     fclose(FILE* stream);
}

int main(int argc, const ubyte** argv)
{
    const ubyte[] text; // never free'd

    if (argc == 2) {
        text = cast(const ubyte[]) readfile(argv[1]);
        if (!text.ptr) {
            fprintf(stderr, "cannot read file\n");
            return 1;
        }
        print_text(text);
    } else if (argc == 3) {
        text = cast(const ubyte[]) readfile(argv[1]);
        if (!text.ptr) {
            fprintf(stderr, "cannot read file\n");
            return 1;
        }
        printf("\n");

        const ubyte[]   buf = cast(const ubyte[]) text;
        const ubyte*    pterm = cast(const ubyte*) argv[2];
        const ubyte[]   term = pterm[0..strlen(pterm)];
        const ubyte[]   func_text;
        int             end;
        int             n;

        while (true) {
            func_text = search(buf, term, &end);

            if (func_text.ptr) {
                n++;
                print_text(func_text);
                printf("\n");
            }

            if (end >= buf.len)
                break;
            buf = buf[end..$];
        }

        if (n == 0) {
            fprintf(stderr, "not found '");
            print_slice(stderr, cast(const ubyte[]) term);
            fprintf(stderr, "'\n");
        }
    } else
        fprintf(stderr, "error: need 0 or 1 arguments\n");

    return 0;
}

void print_text(const ubyte[] s)
{
    if (isatty(STDOUT_FILENO))
        hiprint(s);
    else
        printf("%s", s.ptr);
    printf("\n");
}

// Returns a malloc'd buffer, null-terminated at last position.
//
const ubyte[] readfile(const ubyte* filename)
{
    FILE* f = fopen(filename, "rb");
    if (!f) return null;
    scope fclose(f);

    fseek(f, 0, SEEK_END);
    i64 __size = ftell(f);
    if (__size > 0xffffffff - 1)
        return null;    // file too large
    usz size = cast(usz) __size;
    fseek(f, 0, SEEK_SET);

    ubyte[] buffer = (cast(ubyte*) malloc(size + 1))[0 .. size + 1];
    buffer[size] = 0;

    if (fread(buffer.ptr, 1, size, f) != size) {
        free(buffer.ptr);
        return null;
    }

    return cast(const ubyte[]) buffer;
}

const ubyte[] search(const ubyte[] buf, const ubyte[] term, int* end)
{
    int beg;

    if (!find_function(cast(const ubyte[]) buf, term, &beg, end))
        return null;

    return buf[beg..*end];
}

ubyte getc(const ubyte[] buf, int i)
{
    return i < buf.len ? buf[i] : '\0';
}

void hiprint(const ubyte[] buf)
{
    int i;

    while (true) {
        int a = i;

        switch (getc(buf, i)) {
        case '\0':
            console_reset();
            return;
        case '%':
            i++;
            while (is_ident(getc(buf, i))) i++;
            console_hi(HI.REGISTER);
            break;
        case ';':
            while (getc(buf, i) != '\n' && getc(buf, i) != '\r') i++;
            console_hi(HI.COMMENT);
            break;
        case '@':
            i++;
            while (is_ident(getc(buf, i)) || getc(buf, i) == '.') i++;
            console_hi(HI.ITEM);
            break;
        case '_':
        case 'a': ... case 'z':
        case 'A': ... case 'Z':
            while (is_ident(getc(buf, i))) i++;
            console_hi_word(buf[a..i]);
            break;
        default:
            while (i < buf.len && !is_white(getc(buf, i))) i++;
            while (is_white(getc(buf, i))) i++;
            console_hi(HI.PLAIN);
            break;
        }

        print_slice(stdout, buf[a..i]);
    }
}

bool match(const ubyte[] sa, const ubyte[] sb)
{
    if (sa.len != sb.len)
        return false;

    for (int i; i < sa.len; i++) {
        if (sa[i] != sb[i])
            return false;
    }

    return true;
}

void console_hi_word(const ubyte[] word)
{
    switch (word[0]) {
    case 'a':
        if (match(word, "alloca")) { console_hi(HI.INSTR); return; }
        if (match(word, "alloca")) { console_hi(HI.INSTR); return; }
        if (match(word, "add")) { console_hi(HI.INSTR); return; }
        if (match(word, "and")) { console_hi(HI.INSTR); return; }
        if (match(word, "ashr")) { console_hi(HI.INSTR); return; }
        break;
    case 'b':
        if (match(word, "br")) { console_hi(HI.INSTR); return; }
        if (match(word, "bitcast")) { console_hi(HI.INSTR); return; }
        break;
    case 'c':
        if (match(word, "call")) { console_hi(HI.INSTR); return; }
        break;
    case 'd':
        if (match(word, "define")) { console_hi(HI.ITEM); return; }
        if (match(word, "declare")) { console_hi(HI.ITEM); return; }
        break;
    case 'e':
        if (match(word, "extractvalue")) { console_hi(HI.INSTR); return; }
        break;
    case 'f':
        if (match(word, "fadd")) { console_hi(HI.INSTR); return; }
        if (match(word, "fsub")) { console_hi(HI.INSTR); return; }
        if (match(word, "fmul")) { console_hi(HI.INSTR); return; }
        if (match(word, "fdiv")) { console_hi(HI.INSTR); return; }
        if (match(word, "frem")) { console_hi(HI.INSTR); return; }
        if (match(word, "fpext")) { console_hi(HI.INSTR); return; }
        if (match(word, "fptrunc")) { console_hi(HI.INSTR); return; }
        if (match(word, "fptosi")) { console_hi(HI.INSTR); return; }
        if (match(word, "fptoui")) { console_hi(HI.INSTR); return; }
        if (match(word, "fcmp")) { console_hi(HI.INSTR); return; }
        if (match(word, "fneg")) { console_hi(HI.INSTR); return; }
        break;
    case 'g':
        if (match(word, "getelementptr")) { console_hi(HI.INSTR); return; }
        break;
    case 'i':
        if (match(word, "icmp")) { console_hi(HI.INSTR); return; }
        if (match(word, "inttoptr")) { console_hi(HI.INSTR); return; }
        if (match(word, "insertvalue")) { console_hi(HI.INSTR); return; }
        break;
    case 'l':
        if (match(word, "load")) { console_hi(HI.INSTR); return; }
        if (match(word, "lshr")) { console_hi(HI.INSTR); return; }
        break;
    case 'm':
        if (match(word, "mul")) { console_hi(HI.INSTR); return; }
        break;
    case 'o':
        if (match(word, "or")) { console_hi(HI.INSTR); return; }
        break;
    case 'p':
        if (match(word, "phi")) { console_hi(HI.INSTR); return; }
        if (match(word, "ptrtoint")) { console_hi(HI.INSTR); return; }
        break;
    case 'r':
        if (match(word, "ret")) { console_hi(HI.INSTR); return; }
        break;
    case 's':
        if (match(word, "store")) { console_hi(HI.INSTR); return; }
        if (match(word, "sext")) { console_hi(HI.INSTR); return; }
        if (match(word, "sub")) { console_hi(HI.INSTR); return; }
        if (match(word, "switch")) { console_hi(HI.INSTR); return; }
        if (match(word, "sdiv")) { console_hi(HI.INSTR); return; }
        if (match(word, "srem")) { console_hi(HI.INSTR); return; }
        if (match(word, "shl")) { console_hi(HI.INSTR); return; }
        if (match(word, "sitofp")) { console_hi(HI.INSTR); return; }
        break;
    case 't':
        if (match(word, "trunc")) { console_hi(HI.INSTR); return; }
        break;
    case 'u':
        if (match(word, "udiv")) { console_hi(HI.INSTR); return; }
        if (match(word, "urem")) { console_hi(HI.INSTR); return; }
        if (match(word, "uitofp")) { console_hi(HI.INSTR); return; }
        if (match(word, "unreachable")) { console_hi(HI.INSTR); return; }
        break;
    case 'x':
        if (match(word, "xor")) { console_hi(HI.INSTR); return; }
        break;
    case 'z':
        if (match(word, "zext")) { console_hi(HI.INSTR); return; }
        break;
    default:
        break;
    }

    console_hi(HI.PLAIN);
}

bool find_function(const ubyte[] buf, const ubyte[] term, int* begin, int* end)
{
    int i;
    scope *end = i;

    while (true) {
        while (is_white(buf[i])) i++;

        switch (buf[i++]) {
        case '\0':
            return false;
        case ';':
            while (buf[i] != '\n' && buf[i] != '\r' && buf[i] != '\0') i++;
            i++;
            break;
        case '"':
            while (buf[i] != '"' && buf[i] != '\0') i++;
            i++;
            break;
        case 'd':
            int a = i - 1;
            while (is_ident(buf[i])) i++;

            if (match(buf[a..i], "define")) {
                const ubyte[] name = parse_function(buf, &i);
                if (contain(name, term)) {
                    *begin = a;
                    return true;
                }
            } else if (match(buf[a..i], "declare")) {
                const ubyte[] name = parse_function_signature(buf, &i);
                if (contain(name, term)) {
                    *begin = a;
                    return true;
                }
            }

            break;
        default:
            break;
        }
    }
}

bool contain(const ubyte[] buf, const ubyte[] term)
{
    if (!buf.len || !term.len)
        return false;

    usz b, i, t;

    t = term.len;

    while (i + t <= buf.len) {
        b = i;
        i = i + t;

        while (buf[--i] == term[--t]) {
            if (t == 0)
                return true;
        }

        i = b + 1;
        t = term.len;
    }

    return false;
}

const ubyte[] parse_function_signature(const ubyte[] buf, int* out_i)
{
    const ubyte[] name;

    int i = *out_i;
    scope *out_i = i;

    while (is_white(buf[i])) i++;
    while (buf[i] != '@')    i++;
    i++;

    // function name

    int nb = i;
    int ne = i;

    if (is_ident_start(buf[i])) {
        while (is_ident(buf[i]) || buf[i] == '.') i++;
        ne = i;
    } else if (buf[i] == '\"') {
        while (buf[++i] != '\"') {}
        i++;
        nb++;
        ne = i - 1;
    } else {
        fprintf(stderr, "error: expected function name\n");
        print_current_line(buf, i);
        return null;
    }
    name = buf[nb .. i];

    while (is_white(buf[i])) i++;

    if (!skip_region('(',')', buf, &i))
        return null;

    return cast(const ubyte[]) name;
}

const ubyte[] parse_function(const ubyte[] buf, int* out_i)
{
    int i = *out_i;
    scope *out_i = i;

    const ubyte[] name = parse_function_signature(buf, &i);

    while (is_white(buf[i])) i++;

    while (buf[i] != '{') {
        if (buf[++i] == '\0') {
            fprintf(stderr, "error: no body\n");
            print_current_line(buf, i);
            return null;
        }
    }

    if (!skip_region('{','}', buf, &i))
        return null;

    return cast(const ubyte[]) name;
}

void print_current_line(const ubyte[] buf, int i)
{
    int j = i;
    while (buf[i] != '\n' && buf[i] != '\r' && i > 0) i--;
    int a = i+1;
    while (buf[j] != '\n' && buf[j] != '\r' && j < buf.len) j++;
    fprintf(stderr, "LINE: ");
    print_slice(stderr, cast(const ubyte[]) buf[a..j-1]);
    fprintf(stderr, "\n");
}

void print_slice(FILE* f, const ubyte[] s)
{
    fprintf(f, "%.*s", cast(int) s.len, s.ptr);
}

bool skip_region(ubyte A, ubyte B, const ubyte[] buf, int* out_i)
{
    int i = *out_i;
    scope *out_i = i;

    if (buf[i++] != A) {
        fprintf(stderr, "error: expected '%c'", A);
        return false;
    }

    int nest = 1;

    while (true) {
        ubyte c = buf[i++];
        switch (c) {
        case '\0':
            fprintf(stderr, "error: unclosed parenthesis\n");
            return false;
        case ';':
            while (buf[i] != '\n' && buf[i] != '\r') i++;
            while (buf[i] == '\n' || buf[i] == '\r') i++;
            break;
        default:
            if (c == A) {
                nest++;
            } else if (c == B) {
                nest--;
                if (nest == 0)
                    return true;
            }
            break;
        }
    }
}

bool is_white(ubyte c)
{
    return c == ' ' || c == '\t' || c == '\n' || c == '\r';
}

bool is_alpha(ubyte c)
{
    return c >= 'a' && c <= 'z' ||
           c >= 'A' && c <= 'Z';
}

bool is_num(ubyte c)
{
    return c >= '0' && c <= '9';
}

bool is_ident(ubyte c)
{
    return is_alpha(c) || is_num(c) || c == '_';
}

bool is_ident_start(ubyte c)
{
    return is_alpha(c) || c == '_';
}

//--------------------------------------------------------------------------------

enum HI {
    PLAIN,
    ITEM,
    INSTR,
    REGISTER,
    COMMENT,
}

void console_hi(HI hi)
{
    switch (hi) {
    case HI.PLAIN:
        console_reset();
        break;
    case HI.ITEM:
        console_reset_to_intense();
        break;
    case HI.INSTR:
        console_reset_to_intense();
        break;
    case HI.REGISTER:
        console_reset();
        break;
    case HI.COMMENT:
        console_set(COLOR.lightGray);
        break;
    }
}

/*
 * ANSI escape sequences
 * 0: reset
 * 1: high intensity
 * 2: low intensity
 * 4: underscore
 * 30-37: black, red, .., white
 */

enum COLOR {
    black         = 0,
    red           = 1,
    green         = 2,
    blue          = 4,
    yellow        = red | green,
    magenta       = red | blue,
    cyan          = green | blue,
    lightGray     = red | green | blue,
    bright        = 8,
    darkGray      = bright | black,
    brightRed     = bright | red,
    brightGreen   = bright | green,
    brightBlue    = bright | blue,
    brightYellow  = bright | yellow,
    brightMagenta = bright | magenta,
    brightCyan    = bright | cyan,
    white         = bright | lightGray,
}

void console_reset()
{
    printf("\x1b[m");
}

void console_reset_to_intense()
{
    printf("\x1b[0;1m");
}

void console_set(COLOR c)
{
    printf("\x1b[%d;%dm", (c & COLOR.bright) ? 1 : 0,
                           30 + (c & ~COLOR.bright));
}
